package com.lanyun.aicorrectassistant;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.model.chat.ChatLanguageModel;
import org.apache.poi.xwpf.extractor.XWPFWordExtractor;
import org.apache.poi.xwpf.usermodel.XWPFDocument;

import java.io.FileInputStream;
import java.nio.file.Path;
import java.time.LocalDate;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class SmartHomeworkGrader {
    private static final String SCORE_REGEX = "【总分】\\s*(\\d+)";
    private static final String ACCURACY_REGEX = "【正确率】\\s*([\\d.]+%)";
    private static final String ERROR_ANALYSIS_REGEX = "题号：(\\d+)\\s*错误答案：([A-D])\\s*正确答案：([A-D])\\s*知识点：(.+?)(?=\\n|$)";

    public JSONObject grade(Path filePath, String userPrompt) throws Exception {
        // 读取文档内容
        String content = readDocx(filePath);

        // 生成智能批改指令
        String aiPrompt = buildSmartPrompt(content, userPrompt);

        // 获取AI响应
        String aiResponse = getQwenResponse(aiPrompt);

        // 构建结构化结果
        return buildStructuredResult(content, aiResponse);
    }

    private String readDocx(Path path) throws Exception {
        try (FileInputStream fis = new FileInputStream(path.toFile());
             XWPFDocument doc = new XWPFDocument(fis);
             XWPFWordExtractor extractor = new XWPFWordExtractor(doc)) {
            return extractor.getText();
        }
    }

    private String buildSmartPrompt(String content, String prompt) {
        return String.format(
                "请严格按以下格式分析作业：\n【作业内容】\n%s\n\n【批改要求】\n" +
                        "1. 计算总分（满分100）和正确率\n" +
                        "2. 列出所有错误题目（题号、错误答案、正确答案、知识点）\n" +
                        "3. 总结三大知识薄弱点\n" +
                        "4. 特别关注：%s\n\n" +
                        "请用以下标记分隔：\n【总分】 \n【正确率】 \n【错题分析】 \n【知识总结】",
                content, prompt
        );
    }

    private String getQwenResponse(String prompt) {
        ChatLanguageModel model = QwenChatModel.builder()
                .apiKey("sk-8cc2155ad90049b4bf0ec4bb0ec5cf10")
                .modelName("qwen-max")
                .temperature(0.2F)
                .maxTokens(2000)
                .build();
        return model.chat(prompt);
    }

    private JSONObject buildStructuredResult(String rawContent, String aiResponse) {
        JSONObject result = new JSONObject();

        // 元数据
        Map<String, Object> meta = new HashMap<>();
        meta.put("批改时间", LocalDate.now().toString());
        meta.put("内容摘要", rawContent.substring(0, Math.min(200, rawContent.length())));

        // 核心数据
        Map<String, Object> data = new HashMap<>();
        data.put("score", extractByRegex(aiResponse, SCORE_REGEX, 1));
        data.put("accuracy", extractByRegex(aiResponse, ACCURACY_REGEX, 1));
        data.put("errors", parseErrorAnalysis(aiResponse));
        data.put("weakPoints", parseWeakPoints(aiResponse));

        result.put("metadata", meta);
        result.put("data", data);
        result.put("rawAIResponse", aiResponse);

        return result;
    }

    private List<Map<String, String>> parseErrorAnalysis(String text) {
        List<Map<String, String>> errors = new ArrayList<>();
        Matcher matcher = Pattern.compile(ERROR_ANALYSIS_REGEX, Pattern.DOTALL).matcher(text);

        while (matcher.find()) {
            Map<String, String> error = new HashMap<>();
            error.put("qid", matcher.group(1));
            error.put("studentAnswer", matcher.group(2));
            error.put("correctAnswer", matcher.group(3));
            error.put("knowledgePoint", matcher.group(4).trim());
            errors.add(error);
        }
        return errors;
    }

    private List<String> parseWeakPoints(String text) {
        List<String> points = new ArrayList<>();
        Matcher matcher = Pattern.compile("【知识总结】([\\s\\S]+?)(?=【|$)").matcher(text);

        if (matcher.find()) {
            String[] items = matcher.group(1).split("\\n");
            for (String item : items) {
                if (!item.trim().isEmpty()) points.add(item.trim());
            }
        }
        return points.subList(0, Math.min(3, points.size()));
    }

    private String extractByRegex(String text, String regex, int group) {
        Matcher matcher = Pattern.compile(regex).matcher(text);
        return matcher.find() ? matcher.group(group) : "N/A";
    }

    public static void main(String[] args) throws Exception {
        SmartHomeworkGrader grader = new SmartHomeworkGrader();
        JSONObject result = grader.grade(
                Path.of("C:/Users/Jqcc/Desktop/作业.docx"),
                "请重点检查东北地理常识"
        );
        System.out.println(JSON.toJSONString(result, true));
    }


}